# coding:utf-8
# __user__ = hiicy redldw
# __time__ = 2019/10/8
# __file__ = ONNX_
# __desc__ =
# 一些包的导入
import io
import numpy as np
import torchvision
import torch.utils.model_zoo as model_zoo
import torch.onnx

# 导出onnx 模型
dummy_input = torch.randn(10, 3, 224, 224)
model = torchvision.models.alexnet(pretrained=True)
# 运行一次 跟踪模型运行轨迹
torch.onnx.export(model, (dummy_input,), r"E:\memory\models\alexnet.proto",verbose=True,
                  input_names=['input0'],output_names=['output0'])

"""
简化网络：
    onnx-simplifier
    它推导出整个计算图，然后用其常数输出代替冗余运算符
    pip install onnx-simplifier
    python -m onnxsim input_onnx_model output_onnx_model
"""

# PyTorch中定义的Super Resolution模型
import torch.nn as nn
import torch.nn.init as init
# 超分辨率是一种提高图像，视频分辨率的方法，广泛用于图像处理或视频剪辑。在小型超分辨率模型
class SuperResolutionNet(nn.Module):
    def __init__(self, upscale_factor, inplace=False):
        super(SuperResolutionNet, self).__init__()

        self.relu = nn.ReLU(inplace=inplace)
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 64, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv4 = nn.Conv2d(32, upscale_factor ** 2, (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

        self._initialize_weights()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.pixel_shuffle(self.conv4(x))
        return x

    def _initialize_weights(self):
        init.orthogonal_(self.conv1.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv2.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv3.weight, init.calculate_gain('relu'))
        init.orthogonal_(self.conv4.weight)
# 使用上面模型定义，创建super-resolution模型
torch_model = SuperResolutionNet(upscale_factor=3)
# 加载预先训练好的模型权重
model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'
batch_size = 1    # just a random number

# 使用预训练的权重初始化模型
map_location = lambda storage, loc: storage
if torch.cuda.is_available():
    map_location = None
torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

# 将训练模式设置为falsesince we will only run the forward pass.
torch_model.train(False)

# 输入模型
x = torch.randn(batch_size, 1, 224, 224, requires_grad=True)

# 导出模型
torch_out = torch.onnx._export(torch_model,             # model being run
                               x,                       # model input (or a tuple for multiple inputs)
                               r"E:\memory\models\super_resolution.onnx",  # where to save the model (can be a file or file-like object)
                               export_params=True)      # store the trained parameter weights inside the model file

# 导出给caffe用
import onnx
import caffe2.python.onnx.backend as onnx_caffe2_backend
model = onnx.load(r"E:\memory\models\super_resolution.onnx")
prepared_backend = onnx_caffe2_backend.prepare(model)

W = {model.graph.input[0].name: x.data.numpy()}
# Run the Caffe2 net:
c2_out = prepared_backend.run(W)[0]
# Verify the numerical correctness upto 3 decimal places
np.testing.assert_almost_equal(torch_out.data.cpu().numpy(), c2_out, decimal=3)

print("Exported model has been executed on Caffe2 backend, and the result looks good!")